create_model <- function(docs, seed_list, extra_k){
  set.seed(225)
  names(seed_list) <- 1:length(seed_list)
  dict <- quanteda::dictionary(seed_list)
  model <- topicdict_model(docs,
               dict = dict, extra_k = extra_k,
               remove_numbers = FALSE, 
               remove_punct = TRUE,
               remove_symbols = TRUE,
               remove_separators = TRUE)

  return(model)
}

# Check dispersion
tidy_seededlda_out <- function(model, res, n=15, show=F){
  # Create a nested data frame which contains W and Z
  post <- topicdict::posterior(res)
  topwords <- top_terms(post, n=n)
  topwords <- data.frame(topwords)
  colnames(topwords) <- paste0("EstTopic", 1:ncol(topwords))
  topwords %>%
    tidyr::gather(., key=EstTopic, value=Word) %>%
    mutate(Word = gsub("\\s.*$", "", Word)) -> otidy

  if(show){
    num_seededtopic <- length(model$seeds)
    print(top_terms(post, n)[, 1:num_seededtopic])
  }

  return(otidy)
}


list_to_tibble <- function(lobj){
  # Flatten list and get a tibble
  obj_len <- lobj %>% map(length) %>% flatten_int()
  element <- lobj %>% flatten_chr()

  tibble(SeedTopic = rep(paste0("EstTopic", 1:length(obj_len)), obj_len),
         Word=element
         ) -> res
  return(res)
}

count_appearence_list <- function(otidy, lobj){
  all_words <- lobj %>% flatten_chr()
  SeedTopicName <- paste0("EstTopic", 1:length(lobj))

  otidy %>% 
    right_join(., list_to_tibble(get("lobj")), by="Word") %>%
    mutate(count = ifelse(is.na(EstTopic), 0, 1)) %>%
    group_by(Word, SeedTopic) %>%
    summarize(count = sum(count)) -> organized


  organized %>%
    ggplot(., aes(x=factor(count))) +
    geom_histogram(stat="count") +
    xlab("Number of Topics") + ylab("Count") + ggtitle("Words Split across Topics") +
    theme_bw(base_size=15) +
    theme(plot.title = element_text(hjust = 0.5)) -> g1

  organized %>%
    ggplot(., aes(x=factor(count))) +
    geom_histogram(stat="count") +
    facet_wrap(~ SeedTopic, ncol=3) +
    xlab("Number of Topics") + ylab("Count") + ggtitle("Words Split across Topics") +
    theme_bw(base_size=15) +
    theme(plot.title = element_text(hjust = 0.5)) -> g2

  # Only check topics with seed
  otidy %>% 
    filter(EstTopic %in% get("SeedTopicName")) %>%
    right_join(., list_to_tibble(get("lobj")), by="Word") %>%
    mutate(count = ifelse(is.na(EstTopic), 0, 1)) %>%
    group_by(Word, SeedTopic) %>%
    summarize(count = sum(count)) -> organized_

  organized_ %>%
    ggplot(., aes(x=factor(count))) +
    geom_histogram(stat="count") +
    facet_wrap(~ SeedTopic, ncol=3) +
    xlab("Number of Topics") + ylab("Count") + ggtitle("Words Split across Topics\n(Only topics with keywords)") +
    theme_bw(base_size=15) +
    theme(plot.title = element_text(hjust = 0.5)) -> g3

  return(list(g1, g2, g3))
}

library(tm)

get_lda_result <- function(doc_folder, seed_list, iter_num, k, topicvec=1:k, show_n=15){

  # Prepare Data
  corpus <- Corpus(DirSource(doc_folder))
  strsplit_space_tokenizer <- function(x)
      unlist(strsplit(as.character(x), "[[:space:]]+"))

  dtm <- DocumentTermMatrix(corpus,
                           control = list(tokenize=strsplit_space_tokenizer, 
                           stopwords = F, tolower = T, 
                           stemming = F, wordLengths = c(1, Inf)))

  lda <- LDA(dtm, k = k, control = list(seed = 225, iter=iter_num), method="Gibbs")


  all_words <- seed_list %>% flatten_chr()

  tidytext::tidy(lda) %>%
    group_by(topic) %>%
    top_n(show_n, beta) %>%
    rename(topics = topic, Word = term) %>%
    select(-beta) %>%
    right_join(., list_to_tibble(get("seed_list")), by="Word") %>%
    mutate(count = ifelse(is.na(topics), 0, 1)) %>%
    group_by(Word, SeedTopic) %>%
    summarize(count = sum(count)) -> organized


  organized %>%
    ggplot(., aes(x=factor(count))) +
    geom_histogram(stat="count") +
    xlab("Number of Topics") + ylab("Count") + ggtitle("Words Split across Topics") +
    theme_bw(base_size=15) +
    theme(plot.title = element_text(hjust = 0.5)) -> g1

  organized %>%
    ggplot(., aes(x=factor(count))) +
    geom_histogram(stat="count") +
    facet_wrap(~ SeedTopic, ncol=3) +
    xlab("Number of Topics") + ylab("Count") + ggtitle("Words Split across Topics") +
    theme_bw(base_size=15) +
    theme(plot.title = element_text(hjust = 0.5)) -> g2

  return(list(g1, g2))
}

data_folder <- tempfile()
seed_list <- create_sim_data(saveDir=paste0(data_folder, "Sim1"), D=1000, K=15, TotalV=3000, alpha=0.1, beta_r=0.1, beta_s=0.1, p=c(rep(0.2, 5),rep(0.12, 5),rep(0.05, 5)), lambda=200, seeds_len=5)

seed_list <- lapply(seed_list, function(x){tolower(x)})
seed_list_full <- seed_list
seed_list <- seed_list_full[c(3,4,7,8,14,15)]

doc_folder <- paste0(data_folder, "Sim1", "/W")
docs <- list.files(doc_folder, pattern = "*.txt", full.names = TRUE)
explore_ <- explore(docs,
             remove_numbers = FALSE, # For simulation, make it false
             remove_punct = TRUE,
             remove_symbols = TRUE,
             remove_separators = TRUE)
explore_$visualize_dict_prop(seed_list)

data_folder <- tempfile()
seed_list <- create_sim_data(saveDir=paste0(data_folder, "Sim1"), D=1000, K=45, TotalV=3000, alpha=0.1, beta_r=0.1, beta_s=0.1, p=c(rep(0.2, 15),rep(0.12, 15),rep(0.05, 15)), lambda=200, seeds_len=5)

seed_list <- lapply(seed_list, function(x){tolower(x)})
seed_list_full <- seed_list
seed_list <- seed_list_full[c(5,12,19,24,36,40)]

doc_folder <- paste0(data_folder, "Sim1", "/W")
docs <- list.files(doc_folder, pattern = "*.txt", full.names = TRUE)
explore_ <- explore(docs,
             remove_numbers = FALSE, # For simulation, make it false
             remove_punct = TRUE,
             remove_symbols = TRUE,
             remove_separators = TRUE)
explore_$visualize_dict_prop(seed_list)

diagnosis_topic_recovery_heatmap <- function(post, n=25, title_=T,
                        seed_list=NULL,
                        topicvec=c(), merge=list()){
  topwords <- top_terms(post, n=n)
  topwords <- data.frame(topwords)
  colnames(topwords) <- paste0("EstTopic", 1:ncol(topwords))

  topwords <- tidyr::gather(topwords, key=EstTopic, value=Word) %>%
                mutate(Word = gsub("\\s.*$", "", Word))

  topwords %>%
    mutate(RawWord = Word) %>%
    tidyr::separate(Word,
        into=c("word_id", "TrueTopic"),
        sep="t") %>%
    mutate(TrueTopic = paste0("True", as.character(TrueTopic))) -> res_

  merge_length <- length(merge)
  if(merge_length != 0){
    # Merge Topics
    for(i in 1:merge_length){
      m <- merge[[i]]
      mt <- paste0("True", m)

      res_ %>%
        mutate(TrueTopic=replace(TrueTopic, TrueTopic==mt[1], mt[3])) %>%
        mutate(TrueTopic=replace(TrueTopic, TrueTopic==mt[2], mt[3])) -> res_
    }
  }

  res_ %>%
    group_by(EstTopic, TrueTopic) %>%
    summarise(counts = n()) %>%
    ungroup() %>%
    group_by(EstTopic) %>% 
    mutate(topicsum = sum(counts)) %>%
    ungroup() %>%
    mutate(Proportion = counts / topicsum * 100) -> res_

  if(!is.null(seed_list)){
    # Use only topics with keywords
    num <- length(seed_list)
    seed_list_name <- paste0("EstTopic", 1:num)
    res_ %>%
      filter(EstTopic %in% get("seed_list_name")) -> res_
  }

  num <- length(unique(res_$EstTopic))
  if(is.null(topicvec)){
    res_ %>%
      group_by(EstTopic) %>%
      top_n(1, Proportion) %>%
      mutate(forranking = as.integer(gsub("EstTopic", "", EstTopic))) %>%
      arrange(forranking) %>%
      select(EstTopic) -> topicvec 
    topicvec <-  unique(as.integer(gsub("EstTopic", "", topicvec$EstTopic)))
  }else if(length(topicvec) != num){
    message("topicvec length does not match")
    topicvec <- 1:num
  }

  truenum <- length(unique(res_$TrueTopic))

  title <- paste0("Seeded LDA: Top ", as.character(n), " words")

  g <- ggplot(res_, aes(EstTopic, TrueTopic)) +
        geom_tile(aes(fill=Proportion)) + 
        scale_fill_gradient(limits=c(0, 100), low="#e8e8e8", high="#0072B2", name = "Proportion") +
        scale_x_discrete(limits = rev(paste0("EstTopic", topicvec))) +
        coord_flip() +
        scale_y_discrete(limits = paste0("True", 1:truenum)) +
        xlab("Estimated Topics") + ylab("True Topic") + theme_bw(base_size=13)
  
  if(title_){
    g <- g + ggtitle(title) +
        theme(plot.title = element_text(hjust = 0.5))
  }


  return(g)
}

library(grid)
library(gridExtra)
run_simulations <- function(trueK, estimatedK, seeds_len=6,
                            seed_only=F, seed_contamination=0){
  # Create Combinations
  combinations <- expand.grid(trueK, estimatedK)
  num_combinations <- nrow(combinations)

  # Run Simulations
  for(s in 1:num_combinations){
    trueK_ <- combinations[s, 1]
    estimatedK_ <- combinations[s, 2]

    # Create Data
    set.seed(225)
    data_folder <- tempfile()
    seed_list <- create_sim_data(saveDir=paste0(data_folder, "Sim1"), 
                                 D=1000, K=trueK_, TotalV=3000, alpha=0.1, 
                                 beta_r=0.1, beta_s=0.1, 
                                 p=rep(0.15, trueK_) + rnorm(trueK_, mean=0, sd=0.04), 
                                 lambda=200, seeds_len=seeds_len)
    seed_list <- lapply(seed_list, function(x){tolower(x)})

    # Seed contamination
    if(seed_contamination != 0){
      for(i in 1:seed_contamination){
        seed_list <- lapply(seed_list, function(x){
                      x[sample(1:seeds_len, 1)] <- seed_list[[sample(1:trueK_, 1)]][sample(1:seeds_len, 1)] 
                      return(x)
                      })
      }
    }

    # Fit the model
    extra_k_ <- estimatedK_ - length(seed_list)
    if(extra_k_ < 0){
      message("extra_k is negative, setting it to 0")
      extra_k_ <- 0
      seed_list <- seed_list[1:estimatedK_]
    }
    doc_folder <- paste0(data_folder, "Sim1", "/W")
    docs <- list.files(doc_folder, pattern = "*.txt", full.names = TRUE)
    model <- create_model(docs, seed_list, extra_k=extra_k_)
    res <- topicdict_train(model, iter = iter_num)
    post <- topicdict::posterior(res)

    if(seed_only){
      g <- diagnosis_topic_recovery_heatmap(post, 15, title_=F, seed_list=seed_list)
    }else{
      g <- diagnosis_topic_recovery_heatmap(post, 15, title_=F)
    }

    # Save
    saveRDS(g, file = paste0("/Users/Shusei/Dropbox/Study/Project/ImaiText/topicdict/vignettes/obj/", 
                            "fig_T", trueK_, "_E", estimatedK_, ".obj"))
    saveRDS(post, file = paste0("/Users/Shusei/Dropbox/Study/Project/ImaiText/topicdict/vignettes/obj/", 
                            "post_T", trueK_, "_E", estimatedK_, ".obj"))

    message(paste0("Done: ", s, "/", num_combinations))
  }
}

create_simulation_figure <- function(trueK, estimatedK, title_="Simulation Results"){
  # Create Combinations
  combinations <- expand.grid(trueK, estimatedK) %>%
                    arrange(rev(Var2)) 
  num_combinations <- nrow(combinations)

  # Load Data
  figures <- list()
  for(s in 1:num_combinations){
    trueK_ <- combinations[s, 1]
    estimatedK_ <- combinations[s, 2]

    figures[[s]] <- readRDS(file = paste0("/Users/Shusei/Dropbox/Study/Project/ImaiText/topicdict/vignettes/obj/", 
                            "fig_T", trueK_, "_E", estimatedK_, ".obj"))
  }

  ### Create Figure
  # Get Information
  g1 <- ggplotGrob(figures[[1]])
  id.legend <- grep("guide", g1$layout$name)
  legend <- g1[["grobs"]][[id.legend]]


  # Edit Figure
  edit_figure <- theme(legend.position="none",
                       axis.title.x=element_blank(),
                       axis.title.y=element_blank(),
                       axis.text.x=element_blank(),
                       axis.text.y=element_blank())
  figures <- lapply(figures, function(x){x + edit_figure})

  # New Pictures cf. https://stackoverflow.com/a/11093069/4357279
  g <- arrangeGrob(grobs=figures, 
               nrow = length(estimatedK),
               right = legend,
               top = textGrob(title_),
               left = textGrob("Estimated Topic", rot = 90, vjust = 1),
               bottom = textGrob("True Topic", vjust = -0.1))

  grid.draw(g) # Show plot
}

multiple_simulations <- function(trueK,
                                 estimatedK,
                                 seed_only=F,
                                 seed_contamination=0){

  # Run Simulation
  run_simulations(trueK, estimatedK, seed_only=seed_only, seed_contamination=seed_contamination)

  # Create Figure
  create_simulation_figure(trueK, estimatedK)
}

# How many "true" topic can keywords collect?
multiple_simulations(trueK=c(5,15,25,35), estimatedK=c(5,15,25,35))

# How many "true" topic can keywords collect?
multiple_simulations(trueK=c(5,15,25,35), estimatedK=c(5,15,25,35),
                     seed_contamination=2)